# main.py
import base64
import copy
import json
import logging
import jsonpatch

from flask import Flask, jsonify, Request

app = Flask(__name__)

logging.basicConfig(level=logging.INFO)


def create_patch(modified_object, original_object):
    """创建JSON patch来添加label"""
    if "metadata" not in modified_object:
        modified_object["metadata"] = {}

    if "labels" not in modified_object["metadata"]:
        modified_object["metadata"]["labels"] = {}

    modified_object["metadata"]["labels"]["auto-inject"] = "true"

    return [
        {
            "op": "add",
            "path": "/metadata/labels/auto-inject",
            "value": "true"
        }
    ]


def should_mutate(request_info):
    """检查是否需要进行变更"""
    try:
        # 只处理 default 命名空间
        if request_info["namespace"] != "default":
            return False

        # 只处理 deployment 和 pod
        if request_info["kind"].lower() not in ["deployment", "pod"]:
            return False

        # 检查是否已经有了标签
        if "labels" in request_info["metadata"]:
            if "auto-inject" in request_info["metadata"]["labels"]:
                return False

        return True
    except Exception as e:
        print(f"Error in should_mutate: {str(e)}")
        return False


async def generate_response_body(request_body, env_annotations):
    ingress_rule_name = request_body["request"]["name"]
    object_body = request_body["request"]["object"]
    new_object_body = copy.deepcopy(object_body)

    for key, value in env_annotations.items():
        new_object_body["metadata"]["annotations"][key] = value

    result_generate_response_body = {
        "response": {
            "allowed": True,
            "uid": request_body["request"]["uid"]
        }
    }

    patch_content = jsonpatch.JsonPatch.from_diff(object_body, new_object_body)
    # logging.info(f"patch content:{patch_content}")

    if patch_content:
        result_generate_response_body['response']['patchtype'] = "JSONPatch"
        result_generate_response_body['response']['patch'] = base64.b64encode(str(patch_content).encode()).decode()
        logging.info(f"ingress rule {ingress_rule_name} patch apisix annotations success")
    else:
        logging.info(f"ingress rule {ingress_rule_name} exists apisix annotations,skip")

    return result_generate_response_body


# 添加健康检查端点
@app.route('/mutate', methods=['GET'])
def health_check():
    return jsonify({"status": "healthy"}), 200


@app.route('/mutate', methods=['POST'])
async def mutate(request: Request):
    request_body = await request.json()

    # 输出收到的请求内容
    logging.info("/mutate request: %s", json.dumps(request_body, indent=2))

    # 构造Patch操作，以自动为Pod添加标签
    patch = [
        {
            "op": "add",
            "path": "/metadata/labels/auto-injected",
            "value": "true"
        }
    ]

    # 构造响应体
    response = {
        "apiVersion": "admission.k8s.io/v1",  # 必须包含这个字段
        "kind": "AdmissionReview",  # 必须包含这个字段
        "response": {
            "uid": request_body["request"]["uid"],
            "allowed": True,
            "patchType": "JSONPatch",
            "patch": base64.b64encode(json.dumps(patch).encode()).decode()
        }
    }

    # 返回正确的 JSON 格式以及 HTTP 200 状态码
    return jsonify(response), 200


if __name__ == '__main__':
    # 启动服务器
    # 证书是cert-manager 生成的，自动挂载注入的路径:/certs/tls.crt,/certs/tls.key
    # context = ('/certs/tls.crt', '/certs/tls.key')， 正式环境端口改为443
    context = ('tls.crt', 'tls.key')
    app.run(host='0.0.0.0', port=1443, debug=True, ssl_context=context)
